import typing as tp

import libvirt
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session

from core.configs import db as DB
from core.configs.db import SQLALCHEMY_DATABASE_URL, VIRTUAL_MACHINE_URL
import orator

engine = create_engine(SQLALCHEMY_DATABASE_URL)
MySQLSession = sessionmaker(autocommit=False, autoflush=False, bind=engine)


# Generator[yield_type, send_type, return_type]
def get_db() -> tp.Generator[Session, None, None]:
    db: tp.Optional[Session] = None
    try:
        db = MySQLSession()
        yield db
    finally:
        if db is not None:
            db.close()


def get_vm_conn() -> tp.Generator[libvirt.virConnect, None, None]:
    conn = None
    try:
        conn = libvirt.open(VIRTUAL_MACHINE_URL)
        yield conn
    finally:
        conn.close()


config = {
    'mysql': {
        'driver': 'mysql',
        'host': DB.MYSQL_HOST,
        'database': DB.MYSQL_DB,
        'user': DB.MYSQL_USER,
        'password': DB.MYSQL_PASSWORD,
        'prefix': '',
        'log_queries': True
    }
}

db = orator.DatabaseManager(config)
orator.Model.set_connection_resolver(db)
